"""
方案编写设计和收集者

方案收集者：设计整个图 planner规划(project_name输入) -> writer写作(SubSchemeState输入)->gather收集（SchemeState收集输出）

1.定义writer写作(SubSchemeState输入对象）
2.定义gather收集（SchemeState收集输出对象）
3.定义Designer
1）init方法： llmplanner writer graphe
2）init_graph： 定义三个节点 planner writer gather,
3）定义dispatch节点： plan_node分发write_node, 遍历outline大纲，分发write_node,收集写作内容
4）定义规划节点plan_node，调用规划者
5）定义写作节点write_node，调用写作
6）定义收集节点gather_node，收集写作中contents，放入到 article
4.main:运行设计_call运行图，将article内容结果写入文件中

规划输出数据：
{'abstract': '本项目旨在开发一个基于大规模预训练语言模型的数字人驱动系统，通过集成先进的自然语言处理技术、语音合成和图像生成算法，
实现更加逼真和交互式的数字人。该系统将应用于虚拟客服、娱乐直播等领域，为用户提供沉浸式体验。',
'outline': [
{'name': '项目背景与目标', 'sub_titles': ['行业现状分析', '项目目标设定', '技术挑战与机遇']},
{'name': '关键技术与实现方案', 'sub_titles': ['自然语言处理算法选择', '语音合成系统构建', '图像生成与动画设计']},
{'name': '应用场景与案例研究', 'sub_titles': ['虚拟客服应用示例', '娱乐直播互动场景', '用户体验评估方法']}
]
}

写作：
{'content': '在当前的技术背景下，随着大数据及人工智能技术的迅猛发展')
"""

from langgraph.graph import StateGraph,START,END
from typing import TypedDict,Annotated
from operator import add
from model_utils import getLLM
from plan import create_planner
from writer import create_writer
from langgraph.types import Send

class OutLine(TypedDict):
    title:str
    sub_titles:list[str]

#收集者状态
class SchemeState(TypedDict):
    project_name:str
    abstract:str
    outline:list[OutLine]
    content:Annotated[list[str],add]
    article:str


#writer入口参数
class SubSchemeState(TypedDict):
    project_name:str
    abstract:str
    title:str
    sub_title:str


class Designer:
    def __init__(self):
        self.llm = getLLM()
        self._plan = create_planner(self.llm)
        self._writer = create_writer(self.llm)
        self._graph = self.init_graph()

    def init_graph(self):
        builder = StateGraph(SchemeState)
        builder.add_node("planner_node",self.planner_node)
        builder.add_node("write_node",self.write_node)
        builder.add_node("gather_node",self.gather_node)

        builder.add_edge(START,"planner_node")
        builder.add_conditional_edges("planner_node",self.dispatch_tast)

        builder.add_edge("write_node","gather_node")
        builder.add_edge("gather_node",END)

        return builder.compile()

    def write_node(self,state):
        r = self._writer.invoke(state)
        return {"content":[r["content"]]}

    def planner_node(self,state):
        r = self._plan.invoke(state)
        return {"outline":r["outline"],"abstract":r["abstract"]}

    def gather_node(self,state):
        print("article")
        _text = ""
        _k = 0
        _project_name = state["project_name"]
        _abstract = state["abstract"]
        _outline = state["outline"]
        _content = state["content"]
        _text += _project_name+"\n\n\n"
        _text += _abstract+"\n\n\n"
        for i,title in enumerate(_outline):
            name = title["name"]
            _text += f"{i+1}.{name}\n"
            sub_titles = title["sub_titles"]
            for j,sub in enumerate(sub_titles):
                _text += f"{i+1}.{j}.{sub}\n"
                _text += f"{_content[_k]}\n"
                _k+=1

        return {"article":_text}

    def dispatch_tast(self,state):
        ret = []
        _project_name = state["project_name"]
        _abstract = state["abstract"]
        _outline = state["outline"]
        for title in _outline:
            name = title["name"]
            sub_titles = title["sub_titles"]
            for sub in sub_titles:
                ret.append(Send("write_node",{
                    "project_name":_project_name,
                    "abstract":_abstract,
                    "title":name,
                    "sub_title":sub,
                }))
        print(len(ret))
        return ret

    def __call__(self,project_name):
        r = self._graph.invoke({"project_name": project_name})
        return r

if __name__ == '__main__':
    project_name = "基于大模型的数字人驱动方案"
    des = Designer()
    r = des(project_name)
    with open("article.txt","w+",encoding="utf-8") as f:
        f.write(r["article"])

    print(r["article"])

